import  pandas as pd
import statsmodels.formula.api as smf
from sqlalchemy import create_engine
import pymysql
from matplotlib import pyplot as plt
import seaborn as sns

db_config = {
    'host': 'localhost',
    'user': 'root',
    'password': 'root',
    'database': 'tushare',
    'port': 3306,
    'charset': 'utf8'
}

engine = create_engine(f"mysql+pymysql://{db_config['user']}:{db_config['password']}@{db_config['host']}:{db_config['port']}/{db_config['database']}?charset={db_config['charset']}")

conn = pymysql.connect(**db_config)
chunk_size = 100000

df=pd.read_sql_query("""
                     SELECT  d.*,m.buy_lg_vol,m.sell_lg_vol,m.buy_elg_vol,m.sell_elg_vol,m.net_mf_vol,i.vol as i_vol,i.closes as i_closes FROM date_1 d join
moneyflows m on d.ts_code = m.ts_code and d.trade_date =m.trade_date LEFT JOIN index_daily i on i.trade_date and i.ts_code ='399001.SZ' WHERE d.the_date BETWEEN
'2023-01-01' AND '2023-12.31' AND d.ts_code = '002229.SZ'
                     """
                     ,conn
                     ,chunksize=chunk_size
                     )
df1 = pd.concat(df, ignore_index=True)
print(df1.head())

df1['zd_closes'] = round((df1['closes']  - df1['closes'].shift(1)) / df1['closes'].shift(1),2)
df1['i_closes'] = round((df1['i_closes']  - df1['i_closes'].shift(1)) / df1['i_closes'].shift(1),2)
df1['i_vol'] = round((df1['i_vol']  - df1['i_vol'].shift(1)) / df1['i_vol'].shift(1),2)

df1=df1.dropna(subset=['zd_closes','i_closes','i_vol'])
print(df1.head())

ex=['id','ts_code','trade_date','the_date','opens','high','low','cloese','pre_closes','changes','pct_chg','vol','amount','i_closes']
number=df1.select_dtypes(include=['number']).columns.tolist()
newList = [i for i in number if i not in ex]

formula = 'zd_closes ~ ' + ' + '.join(newList)
res = smf.ols(formula, data=df1).fit()
print(res.summary())    

plt.figure(figsize=(10, 6))
plt.scatter(res.fittedvalues,res.resid)
plt.show()

features=df1[['buy_lg_vol','sell_lg_vol','buy_elg_vol','sell_elg_vol','net_mf_vol','i_vol']]
correlation_matrix = features.corr()

plt.figure(figsize=(10, 6))
sns.heatmap(correlation_matrix, annot=True, cmap='coolwarm', fmt=".2f", linewidths=0.5)
plt.show()